# -*- coding: utf-8 -*-
import tkinter
import traceback

import VBAPItoPython as vb
import os
import xlsxwriter as xl
from PIL import Image
from tkinter import *
from tkinter import filedialog, Listbox, Button, scrolledtext


class application(object):
    def __init__(self):
        self.win = Tk()
        self.inputdir = 'E:/python/test_data-3/'
        self.outputdir = 'E:/python/test_data-3/'
        self.typeflags = {'stp': 0,
                          'asm': 0}
        # self.st_message = scrolledtext.ScrolledText(width=150, height=30)
        # self.st_message.pack()
        self.theLB = Listbox(self.win, selectmode=tkinter.EXTENDED, width=50, height=15)
        self.theLB.grid(row=0, column=0, rowspan=10)
        self.thatLB = Listbox(self.win, selectmode=tkinter.EXTENDED, width=50, height=15)
        self.thatLB.grid(row=0, column=2, rowspan=10)


    def run(self):
        win = self.win
        win.title('批量导出')
        # 创建菜单
        menuBar = Menu(win)

        # 一级菜单 文件
        file_menu = Menu(menuBar, tearoff=0)
        '''
        file_command = [self.file_in, self.file_out, self.set_inputdir, self.set_outputdir, self.file_close]
        file_label = ['导入', '导出', '设置工作目录', '设置输出目录', '关闭']
        for i in range(len(file_label)):
            file_menu.add_command(label=file_label[i], command=file_command[i])
        '''

        file_menu.add_command(label='导入', command=self.file_in)
        # 二级菜单
        export_menu = Menu(file_menu, tearoff=0)
        export_menu.add_command(label='导出文件列表', command=self.export_list)
        export_menu.add_command(label='导出stp', command=self.export_stp)
        file_menu.add_cascade(label='导出', menu=export_menu)

        file_menu.add_command(label='设置工作目录', command=self.set_inputdir)
        file_menu.add_command(label='设置输出目录', command=self.set_outputdir)
        file_menu.add_command(label='关闭', command=self.file_close)

        type_menu = Menu(menuBar, tearoff=0)
        type_menu.add_checkbutton(label='stp', variable=self.typeflags['stp'])
        type_menu.add_checkbutton(label='asm', variable=self.typeflags['asm'])

        help_menu = Menu(menuBar, tearoff=0)
        help_menu.add_command(label='帮助', command=self.help)
        help_menu.add_command(label='关于', command=self.about)

        menuBar.add_cascade(label='文件', menu=file_menu)
        menuBar.add_cascade(label='类型', menu=type_menu)
        menuBar.add_cascade(label='帮助', menu=help_menu)

        win['menu'] = menuBar

        button_addmodel = Button(self.win, text='>>', command=self.addmodel)
        button_addmodel.grid(row=5, column=1)
        button_delmodel = Button(self.win, text='<<', command=self.delmodel)
        button_delmodel.grid(row=6, column=1)
        win.mainloop()

    def file_in(self):
        pass

    def export_list(self):
        # 更改工作目录
        os.chdir(self.outputdir)

        typelist = []
        print(self.typeflags)
        for k in self.typeflags.keys():
            if self.typeflags[k]:
                typelist.append(k)
        print(typelist)
        index = 1
        for file in os.listdir():
            if file.split('.')[-2] in typelist:
                filename = file.split('.')[0]
                data = ('{:<5}'.format(index), '{:<28}'.format(filename))
                self.theLB.insert('end', data)
                # self.st_message.insert('end', data)
                # self.st_message.insert('end', '\n')
                # self.st_message.see(END)
                index += 1

    def addmodel(self):
        while self.theLB.curselection() != ():
            index = self.theLB.curselection()[0]
            self.thatLB.insert('end', self.theLB.get(index))
            self.theLB.delete(index)

    def delmodel(self):
        while self.thatLB.curselection() != ():
            index = self.thatLB.curselection()[0]
            self.theLB.insert('end', self.thatLB.get(index))
            self.thatLB.delete(index)

    def export_stp(self):

        im_width = 60
        im_height = 60
        # 更改工作目录
        os.chdir(self.outputdir)
        # 新建文件
        workbook = xl.Workbook('export.xlsx')
        # 添加工作表
        worksheet = workbook.add_worksheet()
        # 设置单元格格式
        cell_format = workbook.add_format()
        cell_format.set_align('center')
        cell_format.set_align('vcenter')

        # 添加表头
        header = ('No.', '品名/料号', '文件名', '图片', '材质', '外形尺寸')
        worksheet.write_row('A1', header, cell_format)
        # 设置列宽
        worksheet.set_column_pixels(3, 3, im_width)
        worksheet.set_column(1, 2, 23)
        worksheet.set_column(5, 5, 23)
        # 连接现有Creo程序
        connect = vb.IAsyncConnection()
        connect.connect()
        # 数据行计数
        index = 1

        try:
            # 创建会话对象并更改工作目录
            session = vb.Isession(connect.session)
            session.set_directory(Path=self.inputdir)
            # 设置背景为白色 快捷键：wb
            session.runmacro(vb.macro.white_background.value)
            # 遍历会话中所有模型
            for model in session.list_models():
                # 创建模型实例 使用IpfcModel实例化Imodel
                model = vb.Imodel(model)
                if model.extension == 'prt':
                    model.export()  # 输出stp文件
                    model.display()  # 显示模型
                    window = vb.Iwindow(session.get_modelwindow(model.model))  # 创建窗口实例
                    window.activate()  # 激活窗口
                    # 去除模型颜色并居中
                    session.runmacro(vb.macro.clear_appearance.value)
                    session.runmacro(vb.macro.view_refit.value)
                    # 输出模型截图
                    window.export_jpeg(model.instancname + '.jpeg')
                    # 输出模型外形尺寸
                    o = model.get_outline()
                    o.sort(reverse=True)
                    outline = ''
                    for i in o:
                        outline += '{:.3f}'.format(i)
                        outline += '*'
                    outline = outline[:-1]
                    # 生成数据行并插入worksheet
                    data = (index, model.instancname, model.instancname, None, model.material, outline)
                    worksheet.write_row(index, 0, data, cell_format)
                    # 格式化图片并插入表格
                    im = Image.open(model.instancname + '.jpeg')
                    x_scale = (im_width - 10) / im.size[0]
                    y_scale = (im_height - 10) / im.size[1]
                    options = {'x_scale': x_scale,
                               'y_scale': y_scale,
                               'x_offset': 5,
                               'y_offset': 5}
                    worksheet.insert_image(index, 3, model.instancname + '.jpeg', options)
                    # 设置行高
                    worksheet.set_row_pixels(index, im_height)
                    # 数据行计数+1
                    index += 1
                    # 关闭图片 避免后面清理图片时被占用无法清理
                    im.close()
            # 关闭工作部
            workbook.close()
            # 清理无用文件
            for file in os.listdir():
                if not file.split('.')[-1] in ['stp', 'xlsx']:
                    os.remove(file)
        except Exception as e:
            print(traceback.print_exc())
        finally:
            connect.disconnect()

    def set_inputdir(self):
        inputdir = filedialog.askdirectory()
        if inputdir is not None:
            self.inputdir = inputdir
        # self.st_message.insert('end', '工作目录：' + self.inputdir + '\n')
        # self.st_message.see(END)

    def set_outputdir(self):
        outputdir = filedialog.askdirectory()
        if outputdir is not None:
            self.outputdir = outputdir
        # self.st_message.insert('end', '输出目录：' + self.outputdir + '\n')
        # self.st_message.see(END)

    def file_close(self):
        self.win.quit()

    def help(self):
        pass

    def about(self):
        pass


if __name__ == '__main__':
    app = application()
    app.run()
